from pyVim.connect import SmartConnect, Disconnect
from pyVmomi import vim
import ssl
from time import sleep

class VsphereManager:
    def __init__(self, vpshere_server, username, password, dataCenterName, modelFolderName):
        self.vpshere_server = vpshere_server
        self.username = username
        self.password = password
        self.dataCenterName = dataCenterName
        self.modelFolderName = modelFolderName
        self.context = ssl._create_unverified_context()
        self.vm_cache = {}  # 用于缓存虚拟机对象


    def connect(self):
        try:
            self.si = SmartConnect(host=self.vpshere_server, user=self.username, pwd=self.password, sslContext=self.context)
            print("Connected to vSphere API.")
        except Exception as e:
            print(f"Failed to connect to vSphere API: {e}")
            return False
        else:
            self.content = self.si.content
            return True

    def find_datastore(self, find_datacentername):
        datacenteroj = None
        for dc in self.content.rootFolder.childEntity:  
            if hasattr(dc, 'vmFolder') and dc.name == find_datacentername:  
                datacenteroj = dc  
                break
        return datacenteroj

    # 在find_folder方法中，我们通过循环遍历找到目标文件夹，并返回该文件夹的对象。
    def find_folder(self, find_foldername):
        folderoj = None
        datacenter=self.find_datastore(self.dataCenterName)

        if datacenter:  
            folderoj = datacenter.vmFolder  
            folderPath = find_foldername.split('/')  
            for folder in folderPath:  
                if hasattr(folderoj, 'childEntity'):  
                    for child in folderoj.childEntity:  
                        if isinstance(child, vim.Folder) and child.name == folder:  
                            folderoj = child  
                            break  
                    else:  
                        print(f"Folder {folder} not found.")  
                        folderoj = None  
                        break  
        return folderoj

    # 在find_folder_in_root中，我们只在特定数据中心根目录下查找文件夹，不支持递归查找，并返回该文件夹的对象。
    def find_folder_in_root(self, find_foldername):
        print(f"Finding folder {find_foldername} in root folder of {self.dataCenterName}...")
        folderoj = None
        datacenter = self.find_datastore(self.dataCenterName)

        if datacenter:  
            root_folderoj = datacenter.vmFolder
            for child in root_folderoj.childEntity:
                if isinstance(child, vim.Folder) and child.name == find_foldername: 
                    print(f"Found folder: {child.name}")  
                    folderoj = child  
                    break
        return folderoj

    # 在create_folder方法中，在特定数据中心的根目录下创建文件夹，并返回该文件夹的对象。
    def create_folder(self, create_foldername):
        datacenter = self.find_datastore(self.dataCenterName)
        if datacenter:  
            parentFolder = datacenter.vmFolder  # 数据中心的根文件夹 
            folder_obj = parentFolder.CreateFolder(create_foldername)        # 在父文件夹下创建新的文件夹  
            return folder_obj  
        else:  
            print(f"Datacenter {self.dataCenterName} not found.")  
            return None

    # 直接在整个数据中心的下查找虚拟机，缓存虚拟机对象以便后续使用
    def find_vm(self, find_vmname):
         # 如果缓存中存在，直接返回
        if find_vmname in self.vm_cache:
            return self.vm_cache[find_vmname]
        
        # 创建虚拟机视图
        vm_view = self.content.viewManager.CreateContainerView(self.content.rootFolder, [vim.VirtualMachine], True).view
        # 查找虚拟机
        vmNameoj = next((vm for vm in vm_view if vm.name == find_vmname), None)
        # 如果找到，缓存该虚拟机
        if vmNameoj:
            self.vm_cache[find_vmname] = vmNameoj
        return vmNameoj

    # 在给定目录下，来查找指定的虚拟机
    def find_vm_by_folder(self, drictory_name, find_vmname):
        vmNameoj = None
        folder = self.find_folder(drictory_name) 
        if folder:  
            for child in folder.childEntity:  
                if isinstance(child, vim.VirtualMachine) and child.name == find_vmname:  
                    vmNameoj = child  
        else:  
            print("Datacenter not found.") 
        return vmNameoj

    def list_vms_in_folder(self, list_foldername):
        vm_list = [""]
        folder = self.find_folder(list_foldername) 
        if folder:  
            for child in folder.childEntity:  
                if isinstance(child, vim.VirtualMachine):  
                    # print(f"Found VM: {child.name}") 
                    vm_list.append(child.name) 
        else:  
            print("Datacenter not found.") 
        return vm_list

    def power_off_vm(self, vmoj):
        if vmoj.runtime.powerState == vim.VirtualMachinePowerState.poweredOn:
            print(f"正在关闭虚拟机: {vmoj.name}")
            task = vmoj.PowerOff()
            while task.info.state == vim.TaskInfo.State.running:
                pass
            if task.info.state == vim.TaskInfo.State.success:
                print(f"虚拟机{vmoj.name}关闭成功。")
                return True
            else:
                print(f"虚拟机{vmoj.name}关闭失败。")
                return False
        else:
            print(f"虚拟机 {vmoj.name} 已处于关闭状态。")
            return True 
        
    def power_on_vm(self, vmoj):
        if vmoj.runtime.powerState == vim.VirtualMachinePowerState.poweredOff:
            print(f"正在启动虚拟机: {vmoj.name}")
            task = vmoj.PowerOn()
            while task.info.state == vim.TaskInfo.State.running:
                pass
            if task.info.state == vim.TaskInfo.State.success:
                print(f"虚拟机{vmoj.name}启动成功。")
                return True
            else:
                print(f"虚拟机{vmoj.name}启动失败。")
                return False
        else:
            print(f"虚拟机 {vmoj.name} 已处于运行状态。")
            return True 

    def delete_vm(self, vmoj):
        if self.power_off_vm(vmoj):
            print(f"准备删除虚拟机: {vmoj.name}")
            task = vmoj.Destroy()
            while task.info.state == vim.TaskInfo.State.running:
                pass
            if task.info.state == vim.TaskInfo.State.success:
                print(f"虚拟机{vmoj.name}删除成功。")
                return True
            else:
                print(f"虚拟机{vmoj.name}删除失败。")
                return False
        else:
            print(f"虚拟机 {vmoj.name} 关闭失败，无法删除。")
            return False

    def clone_vm(self, source_vmname, target_vmname,create_vm_status):
        vm = self.find_vm(source_vmname)
        datacenter = self.find_datastore(self.dataCenterName)

        if vm:  
            cloneSpec = vim.vm.CloneSpec(powerOn=False) 
            cloneSpec.location = vim.vm.RelocateSpec(datastore=vm.datastore[0]) 

            target_folder = datacenter.vmFolder  
            task = vm.Clone(folder=target_folder, name=target_vmname, spec=cloneSpec)  
            create_vm_status.info(f"Cloning VM {vm.name} to {target_vmname}...")  
            while task.info.state == vim.TaskInfo.State.running:  
                pass
            if task.info.state == vim.TaskInfo.State.success:  
                # print("Clone completed successfully.")  
                create_vm_status.success(f"虚拟机{target_vmname}创建成功!")  
            else:  
                create_vm_status.error(f"虚拟机{target_vmname}创建失败!")
        else:  
            print(f"VM {vm.name} not found.")     

    # 克隆虚拟机到指定文件夹，结合streamlit使用，展示进度状态
    def clone_vm_to_folder_with_streamlit(self, source_vmname, target_vmname, target_foldername, create_vm_status):
        create_vm_status.warning(f"正在查找虚拟机模板对象 {source_vmname}...") 
        vm = self.find_vm(source_vmname)

        if vm:  
            cloneSpec = vim.vm.CloneSpec(powerOn=False) 
            cloneSpec.location = vim.vm.RelocateSpec(datastore=vm.datastore[0]) 

            # 查找并设置目标资源池
            target_pool = self.find_resource_pool(target_foldername)
            print(target_pool)
            if target_pool is None:
                create_vm_status.error(f"资源池 {target_foldername} 不存在！正在创建...")
                target_pool = self.create_resource_pool(target_foldername)
                if target_pool:
                    create_vm_status.success(f"资源池 {target_foldername} 创建成功！")
                    cloneSpec.location.pool = target_pool
                else:
                    create_vm_status.warning(f"资源池 {target_foldername} 创建失败! 将不会指定资源池。")
                sleep(2)
            else:
                cloneSpec.location.pool = target_pool

            # 查找并设置目标文件夹
            target_folder = self.find_folder_in_root(target_foldername)
            if target_folder is None:
                create_vm_status.error(f"组别文件夹 {target_foldername} 不存在！正在创建...")
                target_folder = self.create_folder(target_foldername)

            # 创建并执行克隆任务
            task = vm.Clone(folder=target_folder, name=target_vmname, spec=cloneSpec)  
            create_vm_status.info(f"正在克隆虚拟机模板 {vm.name} 到 {target_vmname}...") 

            while task.info.state == vim.TaskInfo.State.running:  
                progress = task.info.progress
                if progress is not None:
                    create_vm_status.info(f"克隆进度: {progress}%")
                sleep(5)  # 每隔5秒检查一次进度
            if task.info.state == vim.TaskInfo.State.success:
                create_vm_status.success(f"虚拟机{target_vmname}创建成功!")
                return True
            else:  
                create_vm_status.error(f"虚拟机{target_vmname}创建失败!原因: {task.info.error.localizedMessage}")
                return False
        else:  
            create_vm_status.error(f"虚拟机模板 {vm.name} 不存在！请确认虚拟机模板名称是否正确。")
            return False

    # 克隆虚拟机到指定文件夹
    def clone_vm_to_folder(self, source_vmname, target_vmname, target_foldername):
        vm = self.find_vm(source_vmname)

        if vm:  
            cloneSpec = vim.vm.CloneSpec(powerOn=False) 
            cloneSpec.location = vim.vm.RelocateSpec(datastore=vm.datastore[0]) 

            target_folder = self.find_folder_in_root(target_foldername) 
            if target_folder:  
                task = vm.Clone(folder=target_folder, name=target_vmname, spec=cloneSpec)  
                while task.info.state == vim.TaskInfo.State.running:  
                    pass
                if task.info.state == vim.TaskInfo.State.success:  
                    print("Clone completed successfully.")  
                else:  
                    print("Clone failed.") 
            else:
                print(f"Folder {target_foldername} not found.")
        else:  
            print(f"VM {vm.name} not found.")     

    def disconnect(self):
        Disconnect(self.si)

    def add_permission_to_vm(self, vmname, domain, users, role_id):
        vm = self.find_vm(vmname)
        auth_manager = self.content.authorizationManager

        # 创建权限列表
        permissions = []
        for user in users:
            principal = f"{domain}\\{user}"
            permission = vim.Permission()
            permission.principal = principal
            permission.group = False
            permission.propagate = False
            permission.roleId = role_id
            permissions.append(permission)

        # 批量添加权限
        auth_manager.SetEntityPermissions(entity=vm, permission=permissions)

    # 获取权限角色的ID
    def get_role_id(self, role_name):
        role_id = None
        auth_manager = self.content.authorizationManager
        for role in auth_manager.roleList:
            if role.name == role_name:
                role_id = role.roleId
                break
        return role_id

    # 获取已认证的用户
    def get_user_list(self, domain="sznari.com", group_name="svn-研发中心全体员工"):
        users = []
        try:
            search_results = self.content.userDirectory.RetrieveUserGroups(
                searchStr='',
                belongsToGroup=group_name,
                belongsToUser=None,
                domain=domain,
                exactMatch=False,
                findUsers=True,
                findGroups=False,
            )

            for user in search_results:
                if isinstance(user, vim.UserSearchResult):
                    user_name = user.principal if user.principal else '未知'
                    users.append(user_name)
        except Exception as e:
            print(f"获取用户时出现错误: {e}")
        
        return users
    # 资源池处理相关方法
    # 1. 根据名称查找资源池是否存在
    def find_resource_pool(self, pool_name):
        resource_pool = None
        container_view = self.content.viewManager.CreateContainerView(self.content.rootFolder, [vim.ResourcePool], True)
        for managed_object_ref in container_view.view:
            if managed_object_ref.name == pool_name:
                resource_pool = managed_object_ref
                break
        return resource_pool

    # 2. 创建资源池
    def create_resource_pool(self, pool_name):
        datacenter = self.find_datastore(self.dataCenterName)

        if self.find_resource_pool(pool_name):
            print(f"资源池 {pool_name} 已存在！")
            return self.find_resource_pool(pool_name)
        else:
            # 创建 ResourceConfigSpec
            spec = vim.ResourceConfigSpec()

            # 设置 CPU 配置
            cpu_allocation = vim.ResourceAllocationInfo()
            cpu_allocation.reservation = 0  # CPU 预留为 0
            cpu_allocation.limit = -1  # CPU 限制为不受限制
            cpu_allocation.expandableReservation = True  # CPU 分配类型为可扩展
            cpu_allocation.shares = vim.SharesInfo(level=vim.SharesLevel.normal)  # CPU 份额为正常
            spec.cpuAllocation = cpu_allocation

            # 设置内存配置
            mem_allocation = vim.ResourceAllocationInfo()
            mem_allocation.reservation = 0  # 内存预留为 0
            mem_allocation.limit = -1  # 内存限制为不受限制
            mem_allocation.expandableReservation = True  # 内存分配类型为可扩展
            mem_allocation.shares = vim.SharesInfo(level=vim.SharesLevel.normal)  # 内存份额为正常
            spec.memoryAllocation = mem_allocation

            # 遍历集群
            cluster_to_use = None
            for cluster in datacenter.hostFolder.childEntity:
                if isinstance(cluster, vim.ClusterComputeResource):
                    cluster_to_use = cluster
                    break
            if not cluster_to_use:
                return None

            # 创建资源池
            try:
                resource_pool = cluster_to_use.resourcePool.CreateResourcePool(name=pool_name, spec=spec)
            except Exception as e:
                print(f"资源池 {pool_name} 创建失败！{e}")
                return None
            else:
                print(f"资源池 {pool_name} 创建成功！")
                return resource_pool

if __name__ == '__main__':
    # vsphere_server = "10.6.220.37"
    # username = "administrator@vsphere.local"
    # password = "Prs7000.sunri"

    # dataCenterName = "Datacenter"
    modelFolderName = "model"
    vsphere_server = "10.1.109.100"
    username = "yfzxxtb@vsphere.local"
    password = "Qqqq@12345678"

    dataCenterName = "Datacenter"
    modelFolderName = "虚拟机模板/202409"

    vm_manager = VsphereManager(vsphere_server, username, password, dataCenterName, modelFolderName)
    if vm_manager.connect():

        # print(vm_manager.list_vms_in_folder(modelFolderName))
        # xx = vm_manager.find_folder_in_root("model")
        # print(xx)
        # vm_manager.power_off_vm(vm_manager.find_vm("test"))
        # vm_manager.power_on_vm(vm_manager.find_vm("test"))
        # vm_manager.find_vm("CentOS7.3-64（xfce-PRS7072A-PRS-7000S性能测试-吴涛2-16）")
        # vm_manager.clone_vm("test", "test_clone",create_vm_status)
        # vm_manager.clone_vm_to_folder("KylinSec-3-CYSR-240829-60G", "test_clone", "研发测试组")
        # xx =  vm_manager.create_folder("test_folder")
        # vm_manager.add_permission_to_vm("test_clone", "sznari.com", "chenlin1", "普通管理员-CYG")
        # pools = vm_manager.find_vm("厂站系统组","凝思操作系统80-2021-0726（工业集控开发测试-聂军）")
        role_id = vm_manager.get_role_id("普通管理员-CYG")
        print(type(role_id))
        vm_manager.disconnect()